import matplotlib.pyplot as plt

# 定义数据
epochs = list(range(0, 201, 10))

# COD10K 数据集
loss_cod10k_baseline = [0.070, 0.065, 0.060, 0.055, 0.050, 0.045, 0.042, 0.040, 0.038, 0.037, 0.036, 0.035, 0.035, 0.035, 0.035, 0.035, 0.035, 0.035, 0.035]
loss_cod10k_aug = [0.070, 0.065, 0.060, 0.055, 0.050, 0.047, 0.045, 0.042, 0.040, 0.037, 0.035, 0.033, 0.032, 0.032, 0.032, 0.032, 0.032, 0.032, 0.032]
loss_cod10k_fpn = [0.070, 0.065, 0.060, 0.055, 0.050, 0.047, 0.045, 0.042, 0.038, 0.035, 0.032, 0.031, 0.031, 0.031, 0.031, 0.031, 0.031, 0.031, 0.031]
loss_cod10k_aug_fpn_shift_aug_loss = [0.070, 0.065, 0.060, 0.055, 0.050, 0.047, 0.045, 0.042, 0.040, 0.037, 0.034, 0.032, 0.031, 0.030, 0.0296, 0.0295, 0.0295, 0.0294, 0.0293]

# CAMO 数据集
loss_camo_baseline = [0.100, 0.095, 0.090, 0.090, 0.090, 0.088, 0.087, 0.086, 0.085, 0.085, 0.085, 0.085, 0.085, 0.085, 0.085, 0.085, 0.085, 0.085, 0.086]
loss_camo_aug = [0.100, 0.095, 0.090, 0.088, 0.085, 0.083, 0.082, 0.082, 0.082, 0.081, 0.081, 0.081, 0.081, 0.081, 0.081, 0.081, 0.081, 0.081, 0.082]
loss_camo_fpn = [0.100, 0.095, 0.090, 0.088, 0.085, 0.083, 0.082, 0.080, 0.080, 0.079, 0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.080]
loss_camo_aug_fpn_shift_aug_loss = [0.100, 0.095, 0.090, 0.085, 0.080, 0.078, 0.0771, 0.077, 0.077, 0.0769, 0.0766, 0.0765, 0.0764, 0.0763, 0.076, 0.0758, 0.0758, 0.0757, 0.0761]

# CHAMELEON 数据集
loss_chameleon_baseline = [0.050, 0.045, 0.040, 0.035, 0.032, 0.031, 0.030, 0.029, 0.028, 0.027, 0.027, 0.026, 0.026, 0.026, 0.026, 0.026, 0.026, 0.026, 0.031]
loss_chameleon_aug = [0.050, 0.045, 0.040, 0.035, 0.032, 0.031, 0.030, 0.029, 0.028, 0.027, 0.027, 0.026, 0.025, 0.025, 0.025, 0.025, 0.025, 0.025, 0.029]
loss_chameleon_fpn = [0.050, 0.045, 0.040, 0.035, 0.032, 0.031, 0.030, 0.028, 0.027, 0.026, 0.025, 0.024, 0.023, 0.022, 0.022, 0.022, 0.022, 0.021, 0.028]
loss_chameleon_aug_fpn_shift_aug_loss = [0.050, 0.045, 0.040, 0.035, 0.030, 0.028, 0.025, 0.023, 0.023, 0.022, 0.021, 0.020, 0.0208, 0.0207, 0.0207, 0.0205, 0.0205, 0.0204, 0.0204]

# 确保所有数据的长度一致
def adjust_length(data, length):
    return data + [data[-1]] * (length - len(data))

max_len = len(epochs)
loss_cod10k_baseline = adjust_length(loss_cod10k_baseline, max_len)
loss_cod10k_aug = adjust_length(loss_cod10k_aug, max_len)
loss_cod10k_fpn = adjust_length(loss_cod10k_fpn, max_len)
loss_cod10k_aug_fpn_shift_aug_loss = adjust_length(loss_cod10k_aug_fpn_shift_aug_loss, max_len)

loss_camo_baseline = adjust_length(loss_camo_baseline, max_len)
loss_camo_aug = adjust_length(loss_camo_aug, max_len)
loss_camo_fpn = adjust_length(loss_camo_fpn, max_len)
loss_camo_aug_fpn_shift_aug_loss = adjust_length(loss_camo_aug_fpn_shift_aug_loss, max_len)

loss_chameleon_baseline = adjust_length(loss_chameleon_baseline, max_len)
loss_chameleon_aug = adjust_length(loss_chameleon_aug, max_len)
loss_chameleon_fpn = adjust_length(loss_chameleon_fpn, max_len)
loss_chameleon_aug_fpn_shift_aug_loss = adjust_length(loss_chameleon_aug_fpn_shift_aug_loss, max_len)

# 绘制 COD10K 数据集
plt.figure(figsize=(12, 6))
plt.plot(epochs, loss_cod10k_baseline, label='Baseline (UGTR)', marker='o')
plt.plot(epochs, loss_cod10k_aug, label='AUG', marker='o')
plt.plot(epochs, loss_cod10k_fpn, label='AUG+FPN_Shift', marker='o')
plt.plot(epochs, loss_cod10k_aug_fpn_shift_aug_loss, label='AUG+FPN_Shift+AugLoss', marker='o')
plt.title('COD10K Dataset')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig('cod10k_dataset.png')

# 绘制 CAMO 数据集
plt.figure(figsize=(12, 6))
plt.plot(epochs, loss_camo_baseline, label='Baseline (UGTR)', marker='o')
plt.plot(epochs, loss_camo_aug, label='AUG', marker='o')
plt.plot(epochs, loss_camo_fpn, label='AUG+FPN_Shift', marker='o')
plt.plot(epochs, loss_camo_aug_fpn_shift_aug_loss, label='AUG+FPN_Shift+AugLoss', marker='o')
plt.title('CAMO Dataset')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig('camo_dataset.png')

# 绘制 CHAMELEON 数据集
plt.figure(figsize=(12, 6))
plt.plot(epochs, loss_chameleon_baseline, label='Baseline (UGTR)', marker='o')
plt.plot(epochs, loss_chameleon_aug, label='AUG', marker='o')
plt.plot(epochs, loss_chameleon_fpn, label='AUG+FPN_Shift', marker='o')
plt.plot(epochs, loss_chameleon_aug_fpn_shift_aug_loss, label='AUG+FPN_Shift+AugLoss', marker='o')
plt.title('CHAMELEON Dataset')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig('chameleon_dataset.png')